import { post } from "../http";
import { AGENT_BACKEND_API_KEY } from "../auth";

interface EmbeddingRequest {
    input: string[];
    user_id?: string;
}

interface EmbeddingData {
    object: "embedding";
    index: number; // 当前embedding对应的文本在请求中的索引位置
    embedding: number[]; // 文本的向量表示
}

// 定义响应体的接口
interface EmbeddingResponse {
    id: string;
    object: "embedding_list";
    created: number;
    data: EmbeddingData[];
    usage: {
        prompt_tokens: number;
        completion_tokens: number;
        total_tokens: number;
    };
}

/**
 * 获取文本的嵌入向量
 * @param text 文本
 * @returns
 */
export async function getTextEmbedding(
    request: EmbeddingRequest,
    apiKey?: string
): Promise<EmbeddingResponse> {
    const key = apiKey ?? AGENT_BACKEND_API_KEY;
    if (!key) {
        throw new Error("API key is missing");
    }

    const requestBody = {
        ...request
    };

    try {
        const response = await post<EmbeddingResponse>(
            "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/tao_8k?access_token=" +
                key,
            requestBody
        );

        if (response.status === 200) {
            return response.data;
        } else {
            console.error(`Unexpected response status: ${response.status}`);
            throw new Error("Failed to get embedding");
        }
    } catch (error) {
        console.error("Error in getting embedding:", error);
        throw error;
    }
}
